#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
小米便签网页版导出工具（连接已登录浏览器版）

本工具允许用户先在自己的Chrome浏览器中登录小米云服务，然后直接连接到该浏览器实例进行导出。
特点：
- 无需工具自动打开浏览器和处理登录过程
- 支持连接到已有的Chrome浏览器实例
- 只支持JSON和TXT格式导出
- 简单易用，配置灵活
"""

import os
import json
import argparse
import time
import logging
from datetime import datetime
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException, WebDriverException

# 配置日志
def setup_logger():
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler('mi_note_export_existing_browser.log'),
            logging.StreamHandler()
        ]
    )
    return logging.getLogger(__name__)

logger = setup_logger()

class MiNoteWebExporterExistingBrowser:
    def __init__(self, chrome_driver_path, output_dir="./exports", export_format="json", timeout=30):
        """初始化导出工具"""
        self.chrome_driver_path = chrome_driver_path
        self.output_dir = output_dir
        self.export_format = export_format.lower()
        self.timeout = timeout
        self.driver = None
        
        # 验证输出目录存在
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)
            logger.info(f"创建输出目录: {self.output_dir}")
    
    def connect_to_existing_browser(self):
        """连接到已有的Chrome浏览器实例"""
        logger.info("正在准备连接到已有的Chrome浏览器...")
        
        chrome_options = Options()
        chrome_options.add_experimental_option("debuggerAddress", "127.0.0.1:9222")
        
        try:
            # 创建WebDriver服务
            service = Service(executable_path=self.chrome_driver_path)
            
            # 连接到已有的Chrome浏览器实例
            self.driver = webdriver.Chrome(service=service, options=chrome_options)
            logger.info("成功连接到已有的Chrome浏览器实例")
            return True
        except WebDriverException as e:
            logger.error(f"连接Chrome浏览器实例失败: {str(e)}")
            print("\n错误: 连接Chrome浏览器实例失败，请确认：")
            print("1. 是否已使用 --remote-debugging-port=9222 参数启动Chrome浏览器？")
            print("2. ChromeDriver版本是否与Chrome浏览器版本匹配？")
            print("3. 浏览器是否已经打开并可以访问？")
            return False
    
    def check_login_status(self):
        """检查登录状态并导航到便签页面"""
        logger.info("检查浏览器是否已登录小米云服务...")
        try:
            # 导航到小米云服务便签页面
            self.driver.get("https://i.mi.com/note")
            time.sleep(2)  # 等待页面加载
            
            # 检查是否已登录（通过检查便签列表元素是否存在）
            wait = WebDriverWait(self.driver, self.timeout)
            wait.until(EC.presence_of_element_located((By.CLASS_NAME, "note-list-item")))
            
            logger.info("检测到已登录状态，浏览器已成功访问便签页面")
            return True
        except TimeoutException:
            logger.error("未检测到登录状态，无法找到便签列表元素")
            print("\n错误: 未检测到登录状态！请确认：")
            print("1. 您是否已在浏览器中成功登录小米云服务？")
            print("2. 浏览器当前是否显示的是便签页面（https://i.mi.com/note）？")
            print("3. 页面上是否能看到您的便签列表？")
            return False
        except Exception as e:
            logger.error(f"检查登录状态时发生错误: {str(e)}")
            return False
    
    def export_notes(self):
        """导出便签数据"""
        logger.info("开始导出便签数据...")
        print("\n开始导出便签数据，请稍候...")
        
        try:
            # 等待便签列表加载完成
            wait = WebDriverWait(self.driver, self.timeout)
            wait.until(EC.presence_of_element_located((By.CLASS_NAME, "note-list-item")))
            
            # 获取所有便签元素
            note_elements = self.driver.find_elements(By.CLASS_NAME, "note-list-item")
            logger.info(f"找到 {len(note_elements)} 条便签")
            print(f"找到 {len(note_elements)} 条便签，正在逐个导出...")
            
            notes_data = []
            
            # 遍历每个便签并提取内容
            for i, note_element in enumerate(note_elements, 1):
                try:
                    # 点击便签打开详情
                    note_element.click()
                    time.sleep(1)  # 等待内容加载
                    
                    # 提取标题和内容
                    title = "无标题"
                    content = ""
                    
                    # 尝试获取标题（不同版本的UI可能有所不同）
                    try:
                        title_element = self.driver.find_element(By.CLASS_NAME, "note-editor-title")
                        if title_element.text.strip():
                            title = title_element.text.strip()
                    except:
                        try:
                            # 尝试其他可能的标题元素类名
                            title_element = self.driver.find_element(By.CSS_SELECTOR, ".title-input")
                            if title_element.text.strip():
                                title = title_element.text.strip()
                        except:
                            pass
                    
                    # 尝试获取内容
                    try:
                        content_element = self.driver.find_element(By.CLASS_NAME, "note-editor-content")
                        content = content_element.text.strip()
                    except:
                        try:
                            # 尝试其他可能的内容元素类名
                            content_element = self.driver.find_element(By.CSS_SELECTOR, ".content-input")
                            content = content_element.text.strip()
                        except:
                            pass
                    
                    # 提取便签信息
                    note_data = {
                        "id": i,
                        "title": title,
                        "content": content,
                        "export_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    }
                    
                    notes_data.append(note_data)
                    logger.info(f"已提取便签 {i}/{len(note_elements)}: {title}")
                    print(f"已提取便签 {i}/{len(note_elements)}: {title}")
                    
                except Exception as e:
                    logger.error(f"提取便签 {i} 时发生错误: {str(e)}")
                    print(f"警告: 提取便签 {i} 时发生错误，跳过此条")
                    continue
            
            # 保存导出的数据
            self._save_exported_data(notes_data)
            return True
        except Exception as e:
            logger.error(f"导出便签过程中发生错误: {str(e)}")
            return False
    
    def _save_exported_data(self, notes_data):
        """保存导出的数据到文件"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        if self.export_format == "json":
            # 保存为JSON格式
            file_path = os.path.join(self.output_dir, f"mi_notes_{timestamp}.json")
            with open(file_path, "w", encoding="utf-8") as f:
                json.dump(notes_data, f, ensure_ascii=False, indent=2)
            logger.info(f"已将 {len(notes_data)} 条便签导出为JSON格式: {file_path}")
            print(f"\n已将 {len(notes_data)} 条便签导出为JSON格式:")
            print(f"{file_path}")
            
        elif self.export_format == "txt":
            # 保存为TXT格式（每个便签一个文件）
            txt_dir = os.path.join(self.output_dir, f"mi_notes_txt_{timestamp}")
            if not os.path.exists(txt_dir):
                os.makedirs(txt_dir)
            
            for note in notes_data:
                # 生成安全的文件名
                safe_title = "".join([c for c in note["title"] if c.isalnum() or c in " -_"]).rstrip()
                if not safe_title:
                    safe_title = f"无标题_{note['id']}"
                
                file_path = os.path.join(txt_dir, f"{safe_title}.txt")
                
                # 避免文件名冲突
                counter = 1
                base_path = file_path
                while os.path.exists(file_path):
                    name, ext = os.path.splitext(base_path)
                    file_path = f"{name}_{counter}{ext}"
                    counter += 1
                
                # 写入文件
                with open(file_path, "w", encoding="utf-8") as f:
                    if note["title"] != "无标题":
                        f.write(f"{note['title']}\n\n")
                    f.write(note["content"])
                    f.write(f"\n\n导出时间: {note['export_time']}")
            
            logger.info(f"已将 {len(notes_data)} 条便签导出为TXT格式: {txt_dir}")
            print(f"\n已将 {len(notes_data)} 条便签导出为TXT格式（每个便签一个文件）:")
            print(f"{txt_dir}")
            
    def close(self):
        """关闭连接（不会关闭用户的浏览器）"""
        if self.driver:
            self.driver.quit()
            logger.info("已断开与浏览器的连接")

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='小米便签网页版导出工具（连接已登录浏览器版）')
    parser.add_argument('--driver_path', required=True, help='ChromeDriver的路径')
    parser.add_argument('--output_dir', default='./exports', help='导出文件的保存目录')
    parser.add_argument('--format', default='json', choices=['json', 'txt'], help='导出格式')
    parser.add_argument('--timeout', type=int, default=30, help='操作超时时间（秒）')
    return parser.parse_args()

def main():
    print("===== 小米便签网页版导出工具（连接已登录浏览器版）=====")
    print("使用说明：")
    print("1. 请先手动打开Chrome浏览器，并使用以下命令行参数启动：")
    print("   chrome.exe --remote-debugging-port=9222")
    print("2. 在打开的浏览器中登录您的小米云服务账号，并导航到便签页面")
    print("3. 运行本工具，它将自动连接到已登录的浏览器并导出便签")
    print("====================================================")
    
    # 解析命令行参数
    args = parse_args()
    
    logger.info(f"导出格式: {args.format}")
    logger.info(f"输出目录: {args.output_dir}")
    logger.info(f"超时时间: {args.timeout}秒")
    logger.info(f"ChromeDriver路径: {args.driver_path}")
    
    # 验证ChromeDriver是否存在
    if not os.path.exists(args.driver_path):
        logger.error(f"ChromeDriver文件不存在: {args.driver_path}")
        print("错误: ChromeDriver文件不存在，请检查路径是否正确")
        return 1
    
    try:
        # 创建导出工具实例
        exporter = MiNoteWebExporterExistingBrowser(
            chrome_driver_path=args.driver_path,
            output_dir=args.output_dir,
            export_format=args.format,
            timeout=args.timeout
        )
        
        # 连接到已有的Chrome浏览器实例
        if not exporter.connect_to_existing_browser():
            print("\n错误: 连接浏览器失败，请按照上面的使用说明重新操作")
            return 1
        
        # 检查登录状态
        if not exporter.check_login_status():
            print("\n错误: 未检测到登录状态，请先在浏览器中登录小米云服务")
            exporter.close()
            return 1
        
        # 导出便签
        if not exporter.export_notes():
            print("\n错误: 导出便签失败")
            exporter.close()
            return 1
        
        # 完成导出
        exporter.close()
        logger.info("导出完成！")
        print("\n导出成功！")
        print(f"导出文件保存在: {args.output_dir}")
        print("感谢使用小米便签网页版导出工具！")
        return 0
        
    except Exception as e:
        logger.error(f"程序运行失败: {str(e)}")
        print(f"\n错误: {str(e)}")
        return 1

if __name__ == "__main__":
    exit_code = main()
    os._exit(exit_code)